//------------------------------------------------------------------------------
// Includes 
//------------------------------------------------------------------------------
#include <OpenCV/cv.h>
#include <OpenCV/highgui.h>
#include <stdio.h>
#include <ctype.h>
#include <Tracker/CAMShiftTracker.hpp>

namespace Tracker
{
	CAMShiftTracker::CAMShiftTracker() :
	    image(NULL), hsv(NULL), hue(NULL), mask(NULL), backproject(NULL), histimg(NULL),
	    colorBackproject(NULL), hist(NULL), track_object(0), hdims(16), 
	    positionFilter(2), sizeFilter(2), angleFilter(1) {
	
	    setVmin();
	    setVmax();
	    setSmin();
	
	    track_box.center.x = 0;
	    track_box.center.y = 0;
	    track_box.size.width = 0;
	    track_box.size.height = 0;
	    track_box.angle = 0;
	    
	    tracking_color.val[0] = 0;
	    tracking_color.val[1] = 0;
	    tracking_color.val[2] = 0;
	
	    // histogram keeps Hue values from the HSV color space, 
	    // which vary from 0 to 180, and is one dimensional
	    float hranges_arr[] = {0,180};
	    float* hranges = hranges_arr;
	    hist = cvCreateHist(1, &hdims, CV_HIST_ARRAY, &hranges, 1);
	    cvSetZero(hist->bins);
	
	    resetFilters();
	    
	    setMinArea();
	}
	
	CAMShiftTracker::~CAMShiftTracker() {
	    cvReleaseImage(&image);
	    cvReleaseImage(&hsv);
	    cvReleaseImage(&hue);
	    cvReleaseImage(&mask);
	    cvReleaseImage(&backproject);
	    cvReleaseImage(&colorBackproject);    
	    cvReleaseImage(&histimg);
	    cvReleaseHist(&hist);
	}
	
	void CAMShiftTracker::resetFilters() {
	    // the default here assumes an image size of ~320x240 and 
	    // a frame rate of ~30 for a maneuvering time of 1 second
	    positionFilter.setManeuveringTime(30);
	    positionFilter.setProcessStdDev(1);
	    positionFilter.setMeasurementStdDev(2);
	
	    sizeFilter.setManeuveringTime(30);
	    sizeFilter.setProcessStdDev(1);
	    sizeFilter.setMeasurementStdDev(2);
	
	    angleFilter.setManeuveringTime(30);
	    angleFilter.setProcessStdDev(0.01);
	    angleFilter.setMeasurementStdDev(0.05);
	}
	
	void CAMShiftTracker::buildHistogram(CvRect selection) {
	    float max_val = 0.f;
	    cvSetImageROI(hue, selection);
	    cvSetImageROI(mask, selection);
	    cvCalcHist(&hue, hist, 0, mask);
	    cvGetMinMaxHistValue(hist, 0, &max_val, 0, 0);
	    cvConvertScale(hist->bins, hist->bins, max_val ? 255. / max_val : 0., 0);
	    cvResetImageROI(hue);
	    cvResetImageROI(mask);
	}
	
	void CAMShiftTracker::enableTracking(CvRect selection) {
	    // enable the actual tracking...
	    track_window = selection;
	    track_object = 1;
	
	    // we might have a new histogram, update the colorful stuff...
	    if (histimg == NULL) {
	        track_object = -2; // oops, no image, do it later
	        return;
	    }
	    cvZero(histimg);
	    int bin_w = histimg->width / hdims;
	    tracking_color = CV_RGB(0,0,0);
	    for (int i = 0; i < hdims; i++) {
	        int val = cvRound(cvGetReal1D(hist->bins,i)*histimg->height/255);
	        CvScalar color = hsv2rgb(i*180.f/hdims);
	        tracking_color.val[0] += color.val[0]*val;
	        tracking_color.val[1] += color.val[1]*val;
	        tracking_color.val[2] += color.val[2]*val;
	        cvRectangle( histimg, cvPoint(i*bin_w,histimg->height),
	                     cvPoint((i+1)*bin_w,histimg->height - val),
	                     color, -1, 8, 0 );
	    }
	    double maxValue = MAX(MAX(tracking_color.val[0],tracking_color.val[1]),tracking_color.val[2]);
	    tracking_color.val[0] *= 255/maxValue;
	    tracking_color.val[1] *= 255/maxValue;
	    tracking_color.val[2] *= 255/maxValue;
	    
	    double posInit[4] = { selection.x+selection.width/2,  0, 
	                          selection.y+selection.height/2, 0 };
	    positionFilter.init(posInit);
	    double sizeInit[4] = { selection.width,  0, 
	                           selection.height, 0 };
	    sizeFilter.init(sizeInit);
	    double angleInit[2] = { 0, 0 };
	    angleFilter.init(angleInit);
	}
	
	void CAMShiftTracker::run(IplImage *frame) {
	    if (image == NULL) {
	        /* allocate all the buffers */
	        image = cvCreateImage( cvGetSize(frame), 8, 3 );
	        image->origin = frame->origin;
	        hsv = cvCreateImage(cvGetSize(image), 8, 3);
	        hue = cvCreateImage(cvGetSize(image), 8, 1);
	        mask = cvCreateImage(cvGetSize(image), 8, 1);
	        backproject = cvCreateImage(cvGetSize(image), 8, 1);
	        cvZero(backproject);
	        colorBackproject = cvCreateImage(cvGetSize(image), 8, 3);
	        colorBackproject->origin = frame->origin;
	        cvZero(colorBackproject);
	        histimg = cvCreateImage(cvSize(320,200), 8, 3);
	        cvZero(histimg);
	    }
	
	    cvCopy(frame, image, 0);
	    cvCvtColor(image, hsv, CV_BGR2HSV);
	
	    if (track_object) {
	        int _vmin = vmin, _vmax = vmax;
	
	        // cut off pixels whose values are not in the range of interest
	        cvInRangeS(hsv, cvScalar(0,smin,MIN(_vmin,_vmax),0),
	                   cvScalar(180,256,MAX(_vmin,_vmax),0), mask);
	        cvSplit(hsv, hue, 0, 0, 0);
	
	        if (track_object == -1) {
	            // we need to build a histogram from the selection before tracking
	            buildHistogram(selection);
	            enableTracking(selection);
	        } else if (track_object == -2)
	            // just enable tracking, we already have a histogram
	            enableTracking(track_window);
	
	        // make sure the tracking window is within the image or
	        // OpenCV will crash hard.. here, if it is < 0, assume
	        // we want the whole window
	        if (track_window.width < 0)
	            track_window.width = image->width;
	        if (track_window.height < 0)
	            track_window.height = image->height;
	
	        // run the steps required for executing CAMShift...
	        cvCalcBackProject(&hue, backproject, hist);
	        cvAnd(backproject, mask, backproject, 0);
	        cvCamShift(backproject, track_window,
	                   cvTermCriteria( CV_TERMCRIT_EPS | CV_TERMCRIT_ITER, 10, 1 ),
	                   &track_comp, &track_box);
	 
	        // if area is too small, reset from whole frame
	        if (track_comp.area < minArea*255) {
	            track_window.x = track_window.y = 0;
	            track_window.width = backproject->width;
	            track_window.height = backproject->height;
	        } else
	            track_window = track_comp.rect;
	
	        // ??why??
	        track_box.angle = -track_box.angle;
	
	        //
	        // Run the filters to get a better approximation
	        //
	        double positionMeasurement[2] = { track_box.center.x, track_box.center.y };
	        positionFilter.run(positionMeasurement);
	        double positionState[4];
	        positionFilter.getState(positionState);
	        track_box.center.x = positionState[0];
	        track_box.center.y = positionState[2];
	
	        double sizeMeasurement[2] = { track_box.size.width, track_box.size.height };
	        sizeFilter.run(sizeMeasurement);
	        double sizeState[4];
	        sizeFilter.getState(sizeState);
	        track_box.size.width = sizeState[0];
	        track_box.size.height = sizeState[2];
	
	        // we assume the tracker is limited to a range of angle from -Pi/2 to Pi/2
	        // and it is the case with CAMSHIFT, assuming the user does not do 180 degree flips,
	        // if it seems the user switched between the first and second quadrant,
	        // or the third and forth, hack the state before running filter
	        double angleState[2];
	        angleFilter.getState(angleState);
	        double angleMeasurement[1] = { track_box.angle };
	        if (angleState[0] < -CV_PI/4 && angleMeasurement[0] > CV_PI/4) {
	            angleState[0] += CV_PI;
	        } else if(angleState[0] > CV_PI/4 && angleMeasurement[0] < -CV_PI/4) 
	            angleState[0] -= CV_PI;
	        angleFilter.setState(angleState);
	        angleFilter.run(angleMeasurement);
	        angleFilter.getState(angleState);
	        track_box.angle = angleState[0];
	    }
	}
	
	void CAMShiftTracker::drawEllipse(IplImage *image) const { 
	    CvBox2D box = track_box;
	    // make sure the box is valid or OpenCV will crash hard or draw nothing
	    if (box.size.height <= 0)
	        box.size.height = 1;
	    if (box.size.width <= 0)
	        box.size.width = 1;
	
	    if (track_object)
	        cvEllipseBox(image, box, tracking_color, 3, CV_AA, 0);
	}
	
	
	/** 
	 * Gets an RGB color from a Hue value (from HSV color space) 
	 *
	 * @param hue       hue value from 0 to 180
	 * @return          the RGB color in a CvScalar (val[0] = R, etc.) 
	 */
	CvScalar CAMShiftTracker::hsv2rgb(float hue) {
	    int rgb[3], p, sector;
	    static const int sector_data[][3]=
	        {{0,2,1}, {1,2,0}, {1,0,2}, {2,0,1}, {2,1,0}, {0,1,2}};
	    hue *= 0.033333333333333333333333333333333f;
	    sector = cvFloor(hue);
	    p = cvRound(255*(hue - sector));
	    p ^= sector & 1 ? 255 : 0;
	
	    rgb[sector_data[sector][0]] = 255;
	    rgb[sector_data[sector][1]] = 0;
	    rgb[sector_data[sector][2]] = p;
	
	    return cvScalar(rgb[2], rgb[1], rgb[0],0);
	}
}
